"""
Comprehensive field collection implementation for PDEJulia.jl

Based on py-pde field collection system for multi-field coupled PDEs.
Supports heterogeneous field types (scalar, vector, tensor) on the same grid.
"""

using LinearAlgebra
using StaticArrays
using Statistics
using ..Grids

"""
    FieldCollection <: AbstractField

A collection of multiple fields on the same grid, useful for systems of coupled PDEs.

The collection can contain heterogeneous field types (ScalarField, VectorField, Tensor2Field)
as long as they are all defined on the same grid. This is essential for coupled PDE systems
like reaction-diffusion, fluid dynamics, or multi-physics simulations.

# Fields  
- `fields::Vector{AbstractField}`: Vector of fields in the collection
- `labels::Vector{String}`: Optional labels for each field

# Examples
```julia
# Create collection of scalar fields
grid = UnitGrid([64, 64])
u = ScalarField(x -> sin(x[1]), grid)  
v = ScalarField(x -> cos(x[2]), grid)
fields = FieldCollection([u, v], ["u", "v"])

# Mixed field types
density = ScalarField(Float64, grid)
velocity = VectorField(Float64, grid) 
stress = Tensor2Field(Float64, grid)
fields = FieldCollection([density, velocity, stress])

# From function defining all components
fields = FieldCollection(grid, 3) do pos, i
    if i == 1
        return sin(pos[1])  # Scalar
    else
        return [cos(pos[1]), sin(pos[2])]  # Vector (for i > 1)  
    end
end
```
"""
struct FieldCollection <: AbstractField
    fields::Vector{AbstractField}
    labels::Vector{String}
    
    function FieldCollection(fields::Vector{F}, labels::Vector{String} = String[]) where F <: AbstractField
        if isempty(fields)
            throw(ArgumentError("FieldCollection cannot be empty"))
        end
        
        # Check all fields are on the same grid
        ref_grid = grid(fields[1])
        for (i, field) in enumerate(fields[2:end])
            if grid(field) != ref_grid
                throw(ArgumentError("Field $(i+1) has incompatible grid with field 1"))
            end
        end
        
        # Generate default labels if not provided
        if isempty(labels)
            labels = ["field_$i" for i in 1:length(fields)]
        elseif length(labels) != length(fields)
            throw(ArgumentError("Number of labels must match number of fields"))
        end
        
        new(fields, labels)
    end
end

# Convenience constructors
FieldCollection(fields::AbstractField...) = FieldCollection([fields...])

"""
    FieldCollection(field_types::Vector{Type}, grid::AbstractGrid, labels::Vector{String} = String[])

Create collection of zero-initialized fields of specified types.
"""
function FieldCollection(field_types::Vector{Type{F}}, grid::AbstractGrid, 
                        labels::Vector{String} = String[]) where F <: AbstractField
    fields = AbstractField[]
    for field_type in field_types
        if field_type <: ScalarField
            push!(fields, ScalarField(Float64, grid))
        elseif field_type <: VectorField
            push!(fields, VectorField(Float64, grid))
        elseif field_type <: Tensor2Field
            push!(fields, Tensor2Field(Float64, grid))
        else
            throw(ArgumentError("Unsupported field type: $field_type"))
        end
    end
    FieldCollection(fields, labels)
end

"""
    FieldCollection(func::Function, grid::AbstractGrid, n_fields::Int, labels::Vector{String} = String[])

Create collection by evaluating function that returns different field values.
func(pos, field_index) should return the value for field i at position pos.
"""
function FieldCollection(func::Function, grid::AbstractGrid, n_fields::Int, 
                        labels::Vector{String} = String[])
    # Determine field types by sampling the function
    sample_pos = coordinate(grid, (1,) * ndim(grid))
    field_values = [func(sample_pos, i) for i in 1:n_fields]
    
    fields = AbstractField[]
    for (i, val) in enumerate(field_values)
        if val isa Number
            # Scalar field
            field = ScalarField(pos -> func(pos, i), grid)
            push!(fields, field)
        elseif val isa AbstractVector && length(val) == ndim(grid)
            # Vector field  
            field = VectorField(pos -> func(pos, i), grid)
            push!(fields, field)
        elseif val isa AbstractMatrix && size(val) == (ndim(grid), ndim(grid))
            # Tensor field
            field = Tensor2Field(pos -> func(pos, i), grid)
            push!(fields, field)
        else
            throw(ArgumentError("Function returned unsupported type for field $i: $(typeof(val))"))
        end
    end
    
    FieldCollection(fields, labels)
end

# Interface implementation
grid(fc::FieldCollection) = grid(fc.fields[1])

function data(fc::FieldCollection)
    # Return concatenated data vector for all fields
    # This is used for time stepping and linear algebra operations
    all_data = []
    for field in fc.fields
        push!(all_data, vec(data(field)))
    end
    return vcat(all_data...)
end

function Base.similar(fc::FieldCollection)
    similar_fields = [similar(field) for field in fc.fields]
    FieldCollection(similar_fields, copy(fc.labels))
end

function Base.copy(fc::FieldCollection)
    copied_fields = [copy(field) for field in fc.fields]
    FieldCollection(copied_fields, copy(fc.labels))
end

# Indexing and access
Base.length(fc::FieldCollection) = length(fc.fields)
Base.getindex(fc::FieldCollection, i::Int) = fc.fields[i]

function Base.setindex!(fc::FieldCollection, field::AbstractField, i::Int)
    if !(1 ≤ i ≤ length(fc))
        throw(BoundsError(fc, i))
    end
    
    # Check grid compatibility
    if grid(field) != grid(fc)
        throw(ArgumentError("Field has incompatible grid"))
    end
    
    fc.fields[i] = field
end

# Access by label
function Base.getindex(fc::FieldCollection, label::String)
    idx = findfirst(==(label), fc.labels)
    if idx === nothing
        throw(KeyError(label))
    end
    return fc.fields[idx]
end

function Base.setindex!(fc::FieldCollection, field::AbstractField, label::String)
    idx = findfirst(==(label), fc.labels)
    if idx === nothing
        throw(KeyError(label))
    end
    fc[idx] = field
end

# Iterator interface
Base.iterate(fc::FieldCollection) = iterate(fc.fields)
Base.iterate(fc::FieldCollection, state) = iterate(fc.fields, state)

# Get specific field types
function get_scalar_fields(fc::FieldCollection)
    scalar_fields = ScalarField[]
    scalar_indices = Int[]
    
    for (i, field) in enumerate(fc.fields)
        if field isa ScalarField
            push!(scalar_fields, field)
            push!(scalar_indices, i)
        end
    end
    
    return scalar_fields, scalar_indices
end

function get_vector_fields(fc::FieldCollection)
    vector_fields = VectorField[]
    vector_indices = Int[]
    
    for (i, field) in enumerate(fc.fields)
        if field isa VectorField
            push!(vector_fields, field)
            push!(vector_indices, i)
        end
    end
    
    return vector_fields, vector_indices
end

function get_tensor_fields(fc::FieldCollection)
    tensor_fields = Tensor2Field[]
    tensor_indices = Int[]
    
    for (i, field) in enumerate(fc.fields)
        if field isa Tensor2Field
            push!(tensor_fields, field)
            push!(tensor_indices, i)
        end
    end
    
    return tensor_fields, tensor_indices
end# Extended FieldCollection functionality

# Mathematical operations

"""
    +(fc1::FieldCollection, fc2::FieldCollection) -> FieldCollection

Element-wise addition of two field collections.
"""
function Base.:+(fc1::FieldCollection, fc2::FieldCollection)
    if length(fc1) != length(fc2)
        throw(DimensionMismatch("FieldCollections must have same number of fields"))
    end
    
    # Check grid compatibility
    if grid(fc1) != grid(fc2)
        throw(ArgumentError("FieldCollections must have compatible grids"))
    end
    
    # Check field type compatibility
    for i in 1:length(fc1)
        if typeof(fc1[i]) != typeof(fc2[i])
            throw(ArgumentError("Field $i has different type in the two collections"))
        end
    end
    
    new_fields = [fc1[i] + fc2[i] for i in 1:length(fc1)]
    new_labels = [fc1.labels[i] * "+" * fc2.labels[i] for i in 1:length(fc1)]
    
    return FieldCollection(new_fields, new_labels)
end

"""
    -(fc1::FieldCollection, fc2::FieldCollection) -> FieldCollection

Element-wise subtraction of two field collections.
"""
function Base.:-(fc1::FieldCollection, fc2::FieldCollection)
    if length(fc1) != length(fc2)
        throw(DimensionMismatch("FieldCollections must have same number of fields"))
    end
    
    # Check grid compatibility
    if grid(fc1) != grid(fc2)
        throw(ArgumentError("FieldCollections must have compatible grids"))
    end
    
    # Check field type compatibility
    for i in 1:length(fc1)
        if typeof(fc1[i]) != typeof(fc2[i])
            throw(ArgumentError("Field $i has different type in the two collections"))
        end
    end
    
    new_fields = [fc1[i] - fc2[i] for i in 1:length(fc1)]
    new_labels = [fc1.labels[i] * "-" * fc2.labels[i] for i in 1:length(fc1)]
    
    return FieldCollection(new_fields, new_labels)
end

"""
    *(fc::FieldCollection, scalar::Number) -> FieldCollection

Scalar multiplication of field collection.
"""
function Base.:*(fc::FieldCollection, scalar::Number)
    new_fields = [field * scalar for field in fc.fields]
    new_labels = [label * "*$(scalar)" for label in fc.labels]
    return FieldCollection(new_fields, new_labels)
end

Base.:*(scalar::Number, fc::FieldCollection) = fc * scalar

"""
    /(fc::FieldCollection, scalar::Number) -> FieldCollection

Scalar division of field collection.
"""
function Base.:/(fc::FieldCollection, scalar::Number)
    new_fields = [field / scalar for field in fc.fields]
    new_labels = [label * "/$(scalar)" for label in fc.labels]
    return FieldCollection(new_fields, new_labels)
end

"""
    -(fc::FieldCollection) -> FieldCollection

Unary negation of field collection.
"""
function Base.:-(fc::FieldCollection)
    new_fields = [-field for field in fc.fields]
    new_labels = ["-" * label for label in fc.labels]
    return FieldCollection(new_fields, new_labels)
end

# In-place operations
"""
    +!(fc1::FieldCollection, fc2::FieldCollection) -> FieldCollection

In-place addition of field collections.
"""
function Base.:+!(fc1::FieldCollection, fc2::FieldCollection)
    if length(fc1) != length(fc2)
        throw(DimensionMismatch("FieldCollections must have same number of fields"))
    end
    
    for i in 1:length(fc1)
        fc1[i] += fc2[i]
    end
    
    return fc1
end

"""
    -!(fc1::FieldCollection, fc2::FieldCollection) -> FieldCollection

In-place subtraction of field collections.
"""
function Base.:-!(fc1::FieldCollection, fc2::FieldCollection)
    if length(fc1) != length(fc2)
        throw(DimensionMismatch("FieldCollections must have same number of fields"))
    end
    
    for i in 1:length(fc1)
        fc1[i] -= fc2[i]
    end
    
    return fc1
end

"""
    *!(fc::FieldCollection, scalar::Number) -> FieldCollection

In-place scalar multiplication of field collection.
"""
function Base.:*!(fc::FieldCollection, scalar::Number)
    for field in fc.fields
        field *= scalar
    end
    return fc
end

"""
    /!(fc::FieldCollection, scalar::Number) -> FieldCollection

In-place scalar division of field collection.
"""
function Base.:/!(fc::FieldCollection, scalar::Number)
    for field in fc.fields
        field /= scalar
    end
    return fc
end

# Vector operations and norms

"""
    dot(fc1::FieldCollection, fc2::FieldCollection) -> Float64

Compute inner product of two field collections with proper grid weighting.
"""
function LinearAlgebra.dot(fc1::FieldCollection, fc2::FieldCollection)
    if length(fc1) != length(fc2)
        throw(DimensionMismatch("FieldCollections must have same number of fields"))
    end
    
    total_dot = 0.0
    for i in 1:length(fc1)
        total_dot += dot(fc1[i], fc2[i])
    end
    
    return total_dot
end

"""
    norm(fc::FieldCollection, p=2) -> Float64

Compute the p-norm of the field collection.
"""
function LinearAlgebra.norm(fc::FieldCollection, p=2)
    if p == 2
        # L2 norm with proper grid weighting
        total_norm_sq = 0.0
        for field in fc.fields
            total_norm_sq += norm(field, 2)^2
        end
        return sqrt(total_norm_sq)
    elseif p == 1
        # L1 norm
        total_norm = 0.0
        for field in fc.fields
            total_norm += norm(field, 1)
        end
        return total_norm
    elseif isinf(p)
        # L∞ norm
        max_norm = 0.0
        for field in fc.fields
            max_norm = max(max_norm, norm(field, Inf))
        end
        return max_norm
    else
        # General Lp norm
        total_norm_p = 0.0
        for field in fc.fields
            total_norm_p += norm(field, p)^p
        end
        return total_norm_p^(1/p)
    end
end

# Data manipulation and conversion

"""
    to_flat_vector(fc::FieldCollection) -> Vector

Convert field collection to a flat vector for use with linear algebra operations.
"""
function to_flat_vector(fc::FieldCollection)
    return data(fc)
end

"""
    from_flat_vector!(fc::FieldCollection, vector::Vector)

Update field collection from a flat vector.
"""
function from_flat_vector!(fc::FieldCollection, vector::Vector{T}) where T
    offset = 1
    
    for field in fc.fields
        field_size = length(data(field))
        field_data = reshape(vector[offset:offset+field_size-1], size(data(field)))
        
        # Copy data back to field
        data(field) .= field_data
        
        offset += field_size
    end
    
    if offset - 1 != length(vector)
        throw(DimensionMismatch("Vector size doesn't match field collection data size"))
    end
end

"""
    get_field_ranges(fc::FieldCollection) -> Vector{UnitRange{Int}}

Get the index ranges for each field in the flattened data vector.
"""
function get_field_ranges(fc::FieldCollection)
    ranges = UnitRange{Int}[]
    offset = 1
    
    for field in fc.fields
        field_size = length(data(field))
        push!(ranges, offset:offset+field_size-1)
        offset += field_size
    end
    
    return ranges
end

# Field analysis and statistics

"""
    integrate(fc::FieldCollection) -> Vector

Integrate each field in the collection over the domain.
"""
function integrate(fc::FieldCollection)
    results = []
    
    for field in fc.fields
        if field isa ScalarField
            push!(results, integrate(field))
        elseif field isa VectorField
            push!(results, integrate(field))  # Returns vector of integrals
        elseif field isa Tensor2Field
            push!(results, integrate(field))  # Returns matrix of integrals
        end
    end
    
    return results
end

"""
    get_statistics(fc::FieldCollection) -> NamedTuple

Get comprehensive statistics for each field in the collection.
"""
function get_statistics(fc::FieldCollection)
    field_stats = []
    
    for (i, field) in enumerate(fc.fields)
        stats = get_statistics(field)
        push!(field_stats, (
            label = fc.labels[i],
            type = typeof(field),
            stats = stats
        ))
    end
    
    # Collection-wide statistics
    total_norm = norm(fc)
    total_integral = integrate(fc)
    
    return (
        fields = field_stats,
        collection_norm = total_norm,
        collection_integral = total_integral,
        n_fields = length(fc),
        field_types = [typeof(field) for field in fc.fields]
    )
end

"""
    get_extrema(fc::FieldCollection) -> Vector

Get extrema (min, max) for each field in the collection.
"""
function get_extrema(fc::FieldCollection)
    extrema_list = []
    
    for field in fc.fields
        if field isa ScalarField
            push!(extrema_list, get_extrema(field))
        elseif field isa VectorField
            # Get extrema for each component
            stats = get_statistics(field)
            component_extrema = [(comp_stat.min, comp_stat.max) for comp_stat in stats.components]
            magnitude_extrema = (stats.magnitude.min, stats.magnitude.max)
            push!(extrema_list, (components = component_extrema, magnitude = magnitude_extrema))
        elseif field isa Tensor2Field
            # Get extrema for Frobenius norm, trace, determinant
            stats = get_statistics(field)
            tensor_extrema = (
                frobenius = (stats.frobenius_norm.min, stats.frobenius_norm.max),
                trace = (stats.trace.min, stats.trace.max),
                determinant = (stats.determinant.min, stats.determinant.max)
            )
            push!(extrema_list, tensor_extrema)
        end
    end
    
    return extrema_list
end

# Field operations on collections

"""
    apply_to_scalars(fc::FieldCollection, func::Function) -> Vector

Apply function to all scalar fields in the collection.
"""
function apply_to_scalars(fc::FieldCollection, func::Function)
    scalar_fields, indices = get_scalar_fields(fc)
    results = []
    
    for field in scalar_fields
        push!(results, func(field))
    end
    
    return results, indices
end

"""
    apply_to_vectors(fc::FieldCollection, func::Function) -> Vector

Apply function to all vector fields in the collection.
"""
function apply_to_vectors(fc::FieldCollection, func::Function)
    vector_fields, indices = get_vector_fields(fc)
    results = []
    
    for field in vector_fields
        push!(results, func(field))
    end
    
    return results, indices
end

"""
    apply_to_tensors(fc::FieldCollection, func::Function) -> Vector

Apply function to all tensor fields in the collection.
"""
function apply_to_tensors(fc::FieldCollection, func::Function)
    tensor_fields, indices = get_tensor_fields(fc)
    results = []
    
    for field in tensor_fields
        push!(results, func(field))
    end
    
    return results, indices
end

# Specialized constructors for common patterns

"""
    reaction_diffusion_fields(grid::AbstractGrid, n_species::Int) -> FieldCollection

Create field collection for reaction-diffusion system with n chemical species.
Each species is represented as a scalar field.
"""
function reaction_diffusion_fields(grid::AbstractGrid, n_species::Int)
    fields = [ScalarField(Float64, grid) for _ in 1:n_species]
    labels = ["species_$i" for i in 1:n_species]
    return FieldCollection(fields, labels)
end

"""
    fluid_dynamics_fields(grid::AbstractGrid{N}) -> FieldCollection

Create field collection for incompressible fluid dynamics (density + velocity).
"""
function fluid_dynamics_fields(grid::AbstractGrid{N}) where N
    density = ScalarField(Float64, grid)
    velocity = VectorField(Float64, grid)
    
    return FieldCollection([density, velocity], ["density", "velocity"])
end

"""
    elasticity_fields(grid::AbstractGrid{N}) -> FieldCollection

Create field collection for linear elasticity (displacement + stress tensor).
"""
function elasticity_fields(grid::AbstractGrid{N}) where N
    displacement = VectorField(Float64, grid)
    stress = Tensor2Field(Float64, grid)
    
    return FieldCollection([displacement, stress], ["displacement", "stress"])
end

# Boundary condition support

"""
    apply_boundary_conditions!(fc::FieldCollection, boundary_conditions)

Apply boundary conditions to all fields in the collection.
"""
function apply_boundary_conditions!(fc::FieldCollection, boundary_conditions)
    # This would integrate with the boundary condition system
    # For now, it's a placeholder that delegates to each field
    
    if boundary_conditions isa Vector
        # Different BC for each field
        if length(boundary_conditions) != length(fc)
            throw(ArgumentError("Number of boundary conditions must match number of fields"))
        end
        
        for (field, bc) in zip(fc.fields, boundary_conditions)
            apply_boundary_conditions!(field, bc)
        end
    else
        # Same BC for all fields
        for field in fc.fields
            apply_boundary_conditions!(field, boundary_conditions)
        end
    end
    
    return fc
end

# Display and visualization

function Base.show(io::IO, fc::FieldCollection)
    print(io, "FieldCollection(")
    print(io, "$(length(fc)) fields, ")
    print(io, "grid=$(grid(fc))")
    print(io, ")")
end

function Base.show(io::IO, ::MIME"text/plain", fc::FieldCollection)
    println(io, "FieldCollection with $(length(fc)) fields:")
    println(io, "  Grid: $(grid(fc))")
    
    # Show each field with its label and type
    for (i, (field, label)) in enumerate(zip(fc.fields, fc.labels))
        field_type = typeof(field).name.name
        println(io, "  [$i] $label: $field_type")
    end
    
    # Show overall statistics
    stats = get_statistics(fc)
    println(io, "  Collection norm: $(stats.collection_norm)")
    println(io, "  Field types: $(unique(stats.field_types))")
end

# Utility functions

"""
    labels(fc::FieldCollection) -> Vector{String}

Get the labels of all fields in the collection.
"""
labels(fc::FieldCollection) = fc.labels

"""
    relabel!(fc::FieldCollection, new_labels::Vector{String}) -> FieldCollection

Change the labels of the fields in the collection.
"""
function relabel!(fc::FieldCollection, new_labels::Vector{String})
    if length(new_labels) != length(fc)
        throw(ArgumentError("Number of labels must match number of fields"))
    end
    
    fc.labels .= new_labels
    return fc
end

"""
    filter_by_type(fc::FieldCollection, field_type::Type) -> FieldCollection

Create a new collection containing only fields of the specified type.
"""
function filter_by_type(fc::FieldCollection, field_type::Type{T}) where T <: AbstractField
    filtered_fields = AbstractField[]
    filtered_labels = String[]
    
    for (field, label) in zip(fc.fields, fc.labels)
        if field isa T
            push!(filtered_fields, field)
            push!(filtered_labels, label)
        end
    end
    
    return FieldCollection(filtered_fields, filtered_labels)
end

"""
    concatenate(collections::FieldCollection...) -> FieldCollection

Concatenate multiple field collections into one.
All collections must be on the same grid.
"""
function concatenate(collections::FieldCollection...)
    if isempty(collections)
        throw(ArgumentError("Cannot concatenate empty collection list"))
    end
    
    # Check grid compatibility
    ref_grid = grid(collections[1])
    for (i, fc) in enumerate(collections[2:end])
        if grid(fc) != ref_grid
            throw(ArgumentError("Collection $(i+1) has incompatible grid"))
        end
    end
    
    # Combine all fields and labels
    all_fields = AbstractField[]
    all_labels = String[]
    
    for fc in collections
        append!(all_fields, fc.fields)
        append!(all_labels, fc.labels)
    end
    
    return FieldCollection(all_fields, all_labels)
end